package com.example.demo;


import java.util.Date;
import java.util.LinkedList;
import java.util.UUID;

public class SpanManager {


    static ThreadLocal<LinkedList<Span>> spans = ThreadLocal.withInitial(() -> new LinkedList<>());


    static ThreadLocal<String> traceId = ThreadLocal.withInitial(() -> UUID.randomUUID().toString());


    static ThreadLocal<Integer> index = ThreadLocal.withInitial(() -> 0);

    public static Span start(String key) {


        LinkedList<Span> all = spans.get();

        Span span = new Span();
        span.setKey(key);

        Integer dep = index.get();

        span.setSpanId(dep);

        span.setTraceId(traceId.get());
        span.setStartTime(new Date().getTime());

        if (all.size() == 0) {
            span.setParentId(-1);
        } else {
            span.setParentId(all.peekLast().getSpanId());
        }
        all.add(span);
        index.set(index.get() + 1);

        return span;
    }


    public static void stop(Span span) {

        if (span == null) {
            return;
        }

        LinkedList<Span> all = spans.get();

        if (all.size() == 0) {
            return;
        }

        Span last = all.peekLast();

        if (!last.getTraceId().equals(span.getTraceId())
                || last.getSpanId() != span.getSpanId()) {
            return;
        }

        Span remove = all.removeLast();
        index.set(index.get() - 1);

        remove.setEndTime(new Date().getTime());

        remove.setCostMs((int) (remove.getEndTime() - remove.getStartTime()));

        System.out.println(remove);

        if (spans.get().size() == 0) {
            spans.remove();
            traceId.remove();
            index.remove();
        }

    }


}
